from msys_fig import PlotSurf
from numpy import array, linspace, meshgrid, zeros
from pylab import contour, show, axis, axhline, axvline, plot, grid, xlabel, ylabel, subplot, text
import sys

def find_first_root(x0, f, zet, tolF=1e-15, tolDx=1e-15, maxIt=20):
    x = x0
    print '[1;32m%6s%14s%14s%14s%14s[0m' % ('it', 'x', 'fx', 'z', 'N')
    for it in range(maxIt):
        fx = f(x)
        if abs(fx) < tolF: break
        fm = f(x - zet*fx)
        fp = f(x + zet*fx)
        z  = fx - fm
        N  = z*z/(z*z + fx*(z + fx - fp))
        print '%6d%14.9f%14.9f%14.9f%14.9f' % (it, x, fx, z, N)
        dx = zet*abs(N)*fx*fx/abs(z)
        if abs(dx) < tolDx: break
        x += dx
    else: print '[1;31mfind_first_root failed with x=%g[0m' % x
    return x

# triangle
#def ftri(r, s, t): return 27.0*r*s*(1.0-r-s)
# quad
def fqua(r, s, t): return -(1.0-r**2.0)*(1.0-s**2.0)
# tet
#def ftet(r, s, t): return 256.0*r*s*t*(1.0-r-s-t)
# hex
#def fhex(r, s, t): return (1.0-r**2.0)*(1.0-s**2.0)*(1.0-t**2.0)

t3d = False
nxx = 41
if t3d: nxx = 21

X = linspace(-1,1,nxx)
Y = linspace(-1,1,nxx)
XX, YY = meshgrid(X,Y)
Z = fqua(XX, YY, 0)
if t3d:
    PlotSurf(XX, YY, Z)
    show()
    sys.exit(0)
else:
    subplot(2,1,1)
    p = contour(XX, YY, Z)
    grid(color='grey')
    plot([-1, 1], [-1,-1], 'k-', clip_on=0)
    plot([-1, 1], [ 1, 1], 'k-', clip_on=0)
    plot([ 1, 1], [-1, 1], 'k-', clip_on=0)
    plot([-1,-1], [-1, 1], 'k-', clip_on=0)

u0 = array([-1., -0.2])
uu = array([2.4, 1.6])
def falp(alp, doplot=False):
    u = u0 + alp * uu
    plot(u[0], u[1], 'g.')
    return fqua(u[0], u[1], 0)
zet = 0.5
alp_root = find_first_root(1e-3, falp, zet, maxIt=20)
print 'alp_root =', alp_root

np    = 21
dalp  = 1.0
A, Fa = zeros(np), zeros(np)
for i in range(np):
    alp = i*dalp/float(np-1)
    fa  = falp(alp, doplot=False)
    print 'falp(%g) = %g' % (alp, fa)
    A[i], Fa[i] = alp, fa

u  = u0 + uu
ua = u0 + alp_root*uu
if not t3d:
    plot([u0[0],u[0]], [u0[1],u[1]], 'k-')
    plot(u0[0], u0[1], 'ro')
    plot(ua[0], ua[1], 'r^')
    axis('equal')

subplot(2,1,2)
plot(A,Fa)
xlabel('alpha'); ylabel('f(alpha)')
axhline(0,color='k',ls='--')
axvline(alp_root,color='k',ls='--')
text(alp_root,0.,'%g'%alp_root)
grid(color='grey')

show()
